/*
  Copyright 2006-2008 Patrik Jonsson, sunrise@familjenjonsson.org

  This file is part of Sunrise.

  Sunrise is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  Sunrise is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.
  
*/

#include <iostream>
#include "CCfits/CCfits"
#include "counter.h"
#include "polychromatic_grid.h"
#include "optical.h"
#include "scatterer.h"
#include "constants.h"
#include "dust_model.h"
#include "emission.h"
#include "emission-fits.h"
#include "xfer.h"
//#include "single_lambda_emergence.h"
#include "full_sed_emergence.h"
#include "full_sed_grid.h"
#include "aux_grid.h"
#include "aux_emergence.h"
#include "emergence-fits.h"
#include "mcrx-units.h"
#include "grid_factory.h"
#include "grid-fits.h"
#include "shoot.h"
#include "fits-utilities.h"
#include "boost/lexical_cast.hpp"
#include "density_generator.h"

//using namespace blitz;
using namespace mcrx;
using namespace CCfits;
using namespace std; 
class dummy_grid {
public:
  typedef void T_cell;
};

typedef local_random T_rng_policy;
typedef scatterer<polychromatic_scatterer_policy, T_rng_policy> T_scatterer;
typedef T_scatterer::T_lambda T_lambda;
typedef T_scatterer::T_biaser T_biaser;
typedef dust_model<T_scatterer, cumulative_sampling, T_rng_policy> T_dust_model;
typedef full_sed_emergence T_emergence;

typedef full_sed_grid T_grid;

const int n_species = 2;


class uniform_factory: public grid_factory<T_grid::T_data> {
public:
  const T_densities rho;
  const int n_lambda, ref;

  uniform_factory (const T_densities& r, int nl, int rl):
    rho (r), n_lambda (nl), ref(rl) {};
  
  virtual bool refine_cell_p (const T_cell& c, int level) {
    return level <ref;};
  virtual bool unify_cell_p (T_grid::const_local_iterator begin,
			     T_grid::const_local_iterator end) {
    return false;
  };
  virtual T_data get_data (const T_cell& c) {
    return T_data(rho, vec3d(0,0,0), T_lambda(n_lambda));
  }
    
  virtual int n_threads () {return 1;};
  virtual int work_chunk_levels () {return 5;};
  virtual int estimated_levels_needed () {return 0;};
};


class w77_factory: public grid_factory<T_grid::T_data> {
public:
  const T_float radius, thickness, tau, z0;
  const int n_lambda;

  w77_factory (T_float r, T_float t, T_float zz, T_float o, int nl):
    radius (r), z0 (zz), thickness (t), tau  (o), n_lambda (nl) {};
  
  virtual bool refine_cell_p (const T_cell& c, int level) {
    return level <4;};
  virtual bool unify_cell_p (T_grid::const_local_iterator begin,
			     T_grid::const_local_iterator end) {
    T_densities o = (begin++)->data()->get_absorber().densities();
    while (begin != end)
      if (any(o !=(begin++)->data()->get_absorber().densities()))
	return false; 
    return true;
  };
  virtual T_data get_data (const T_cell& c) {
    const vec3d center = c.getcenter();
    T_densities opacity(n_species);
    opacity=0;
    const T_float r2 = center [0]*center  [0] + center [1]*center [1]; 
    if ((r2 < radius*radius) && 
	(center [2] < -z0) && (center [2] > -z0-thickness))
      opacity = tau/thickness;
    return T_data (opacity, vec3d(0,0,0), T_lambda(n_lambda));
  }
    
  virtual int n_threads () {return 1;};
  virtual int work_chunk_levels () {return 5;};
  virtual int estimated_levels_needed () {return 0;};
};




void w77run (T_float tau, T_float radius, 
	     T_float thickness, T_float z0,
	     T_float g, T_float a,
	     int n_forced, T_float distance,
	     const string& name,
	     bool integrated = false,
	     bool collimated = false)
{
  cout << "loading grid" << endl;
  const T_float grid_extent = 1.1*radius;
  const T_float image_extent = 1.2*grid_extent;
  const long n_rays =1000000;
  const int grid_n = 3;
  const int npix = 101;
  std::vector<std::pair<T_float, T_float> > cam_pos;
  cam_pos.push_back(std:: make_pair (10.*constants::pi/180, 0.) );
  /*
  cam_pos.push_back(std:: make_pair (0.*pi/180, 0.) );
  cam_pos.push_back(std:: make_pair (20.*pi/180, 0.) );
  cam_pos.push_back(std:: make_pair (30.*pi/180, 0.) );
  cam_pos.push_back(std:: make_pair (40.*pi/180, 0.) );
  cam_pos.push_back(std:: make_pair (50.*pi/180, 0.) );
  cam_pos.push_back(std:: make_pair (60.*pi/180, 0.) );
  cam_pos.push_back(std:: make_pair (70.*pi/180, 0.) );
  cam_pos.push_back(std:: make_pair (80.*pi/180, 0.) );
  cam_pos.push_back(std:: make_pair (90.*pi/180, 0.) );
  cam_pos.push_back(std:: make_pair (100.*pi/180., 0.) );
  cam_pos.push_back(std:: make_pair (110.*pi/180., 0.) );
  cam_pos.push_back(std:: make_pair (120.*pi/180., 0.) );
  cam_pos.push_back(std:: make_pair (130.*pi/180., 0.) );
  cam_pos.push_back(std:: make_pair (140.*pi/180., 0.) );
  cam_pos.push_back(std:: make_pair (150.*pi/180., 0.) );
  cam_pos.push_back(std:: make_pair (160.*pi/180., 0.) );
  cam_pos.push_back(std:: make_pair (170.*pi/180., 0.) );
  cam_pos.push_back(std:: make_pair (180.*pi/180., 0.) );
  */
  // k g a
  T_dust_model::T_scatterer_vector sv;
  const int n=13;
  T_lambda kk(n),kk1(n),kk2(n),gg(n),aa(n);
  kk=3.,3.,3.,3.,3.,1.,1.,1.,1.,1.,2.,4.,5.;
  kk1=.3,.3,.3,.3,.3,.1, .9,.9,.9,.9,1.8,3.6,4.5;
  kk2=2.7,2.7,2.7,2.7,2.7,.9, .1,.1,.1,.1,.2,.4,.5;
  gg=.1,.3,.5,.7,.9,.7,.7,.7,.7,.7,.7,.7,.7;
  aa=.6,.6,.6,.6,.6,.2,.4,.6,.8,1.,.6,.6,.6;
  //kk1=0.5;gg=g;aa=a;
  //sv.push_back(boost::shared_ptr<T_scatterer> (new simple_HG_dust_grain<polychromatic_scatterer_policy>(kk,gg,aa)));
  sv.push_back(boost::shared_ptr<T_scatterer> (new simple_HG_dust_grain<polychromatic_scatterer_policy, T_rng_policy>(kk1,gg,aa)));
  sv.push_back(boost::shared_ptr<T_scatterer> (new simple_HG_dust_grain<polychromatic_scatterer_policy, T_rng_policy>(kk2,gg,aa)));

  //for(int i=0;i<n_species-1;++i)
  //sv.push_back(boost::shared_ptr<T_scatterer> (new simple_HG_dust_grain<polychromatic_scatterer_policy>(kk,gg,aa)));
  //sv.push_back(boost::shared_ptr<T_scatterer> (new simple_HG_dust_grain<polychromatic_scatterer_policy>(kk,gg,aa)));
  T_dust_model model (sv);
  T_lambda lambda(n);
  lambda = 550e-9;
  model.set_wavelength(lambda);

  const T_lambda albedo = model.get_scatterer(0 ).albedo();

  w77_factory factory (radius, thickness, z0, tau, n);
  polychromatic_grid gr (vec3d (-grid_extent,-grid_extent,-z0-thickness ), 
	     vec3d (grid_extent,grid_extent,-z0 ),
	     vec3i (grid_n, grid_n, 1), factory);

  cout << "setting up emission" << endl;

  typedef emission<polychromatic_policy, T_rng_policy> T_emission;
  auto_ptr<T_emission> em;
  //if (collimated)
  //em.reset(new laser_emission<T_lambda, T_rng_policy> (vec3d (1e-20, 1e-20, 1e-20), vec3d (0, 0, -1), 1.));
  //else
  T_lambda source(n);
  source = 1;
  em.reset(new pointsource_emission<polychromatic_policy, T_rng_policy> (vec3d (1e-10, 1e-10, 1e-10), source));
  
  cout << "setting up emergence" << endl;
  
  //auto_ptr<scatter_emergence>  e (new scatter_emergence (cam_pos, distance, image_extent, npix));
  T_emergence e (cam_pos, distance, image_extent, npix);
  
  cout << "shooting" << endl;

  const int n_threads = 1;
  std::vector<T_rng::T_state> states= generate_random_states (n_threads);
  T_rng_policy rng;

  xfer<T_dust_model, polychromatic_grid, T_rng_policy> x (gr, *em, e, model, rng, 1e-4, 0,n_forced);
  long n_rays_actual = 0;
  dummy_terminator t;
  T_biaser b(0);
  //shoot (x, scatter_shooter<T_biaser> (b), t, states, n_rays,
  //	 n_rays_actual, n_threads,0, "");

  for(;n_rays_actual<n_rays;n_rays_actual++) {
    T_biaser b(int(rnd()*n));
    //T_biaser b(7);
    x.shoot(b);
  }


  T_unit_map units;
  units ["length"] = "whatever";
  auto_ptr<FITS>  output;
  output.reset (new FITS ("!"  + name + ".fits", Write));
  //e.write_parameters(output, units);
  //e.create_HDU(*output, 2, units);

  cout << "saving" << endl;
  //e.create_HDU(output, 1, g.units);

  const T_float normalization = 1.//em->total_emission_weight()
    /n_rays_actual;
  //global_hdu = output.addImage ("junk", DOUBLE_IMG, n );
  //e.write_images(output, emission.total_emission_weight()/N, 1);

  //FITS output ("junk.fits", Write);

  if (!integrated) {
    // extract data and save a text file
    ofstream of ((name + ".txt").c_str());
    for (T_emergence::iterator i = e.begin(); i != e.end(); ++i) {
      const T_emergence::T_camera::T_image image =
	(*i)->get_image ();
      const int npix = image.extent(0 );
      const T_float scale = image_extent/distance*206265/npix ; // arc second per pixel
      for (int j = 0; j < npix; ++j) {
	of << (j- (npix-1)/2)*scale;
        for(int k=0; k<image.extent(thirdDim);++k)
	  of << '\t' <<image ((npix-1)/2,j,k)*normalization/(*i)->pixel_solid_angle();
	of << '\n';
      }
      of << "\n\n";
    }
  }
  else {
    // extract data and save a text file
    ofstream of ((name + ".txt").c_str());
    for (T_emergence::iterator i = e.begin(); i != e.end(); ++i) {
      const T_emergence::T_camera::T_image image =
	(*i)->get_image ();
      of << (*i)->get_theta() << '\t' 
	 << sum (image)*normalization*distance*distance << '\n';
    }
  }

  {
    output  .reset (new FITS (name + ".fits",  Write  ));
    //e.write_images(*output,normalization, 1);
    e.write_images(*output,normalization, units, "", false, false);
  }

}

void w77 ()
{
  cout << "Testing results from Witt 1977\n";
  
  //w77run (1, 1.2, 1,.05, .01, .6, 1, 126, "poly");
  //w77run (5, 1.2, 1,.05, .01, .6, 2, 126, "w77t12_tau5");
  //w77run (1, 1.2, 1,.05, .7, .6, 2, 126, "w77t34_tau1");
  //w77run (5, 1.2, 1,.05, .7, .6, 2, 126, "w77t34_tau5");

  // now paper III
  //w77run (3, 1.2, 1,-.5, .1, .6, 2, 126, "w77III_g.1_tau3_a.6");
  //w77run (3, 1.2, 1,-.5, .3, .6, 2, 126, "w77III_g.3_tau3_a.6");
  //w77run (3, 1.2, 1,-.5, .5, .6, 2, 126, "w77III_g.5_tau3_a.6");
  //w77run (3, 1.2, 1,-.5, .7, .6, 2, 126, "w77III_g.7_tau3_a.6");
  //w77run (3, 1.2, 1,-.5, .9, .6, 2, 126, "w77III_g.9_tau3_a.6");
  //w77run (3, 1.2, 1,-.5, .9, .6, 1, 126, "w77III_g.9_tau3_a.6_nf1");

  w77run (1, 1.2, 1,-.5, .7, .2, 1, 126, "poly2");
  //w77run (1, 1.2, 1,-.5, .7, .2, 2, 126, "w77III_g.7_tau1_a.2");
  //w77run (1, 1.2, 1,-.5, .7, .4, 2, 126, "w77III_g.7_tau1_a.4");
  //w77run (1, 1.2, 1,-.5, .7, .6, 2, 126, "w77III_g.7_tau1_a.6");
  //w77run (1, 1.2, 1,-.5, .7, .8, 2, 126, "w77III_g.7_tau1_a.8");
  //w77run (1, 1.2, 1,-.5, .7, 1., 2, 126, "w77III_g.7_tau1_a1");
  //w77run (1, 1.2, 1,-.5, .7, 1., 1, 126, "w77III_g.7_tau1_a1_nf1");

  // this is very similar to the test by Watson and Henney
  //w77run (2, 10000, 1,0, .5, .5, 0, 1000000, "wh01a",true);
  //w77run (2, 10000, 1,0, .5, .5, 3, 1000000, "wh01b",true, true);
  
}

void galaxy ()
{
  cout << "loading grid" << endl;

  FITS file("/data/sunrise/test/Sbc201a-u4/timetest/grid_012.fits",Read);
  ExtHDU& lambda_hdu=open_HDU(file,"LAMBDA");
  T_lambda lambda;
  read(lambda_hdu.column("lambda"),lambda);
   const int n_lambda = lambda.size();

   T_dust_model::T_scatterer_vector sv;
   sv.push_back(boost::shared_ptr<T_scatterer> 
		(new Draine_grain<polychromatic_scatterer_policy, T_rng_policy>("/home/sunrise/dust_data/kext_albedo_WD_MW_3.1_60")));

  T_dust_model model (sv);
  model.set_wavelength(lambda);
  uniform_density_generator dg(0,0.4);

  T_grid gr(open_HDU(file,"GRIDSTRUCTURE"),open_HDU(file,"GRIDDATA"),
	    dg, n_lambda);

  cout << "setting up emission" << endl;
  typedef full_sed_particles_emission<cumulative_sampling, T_rng_policy> T_emission;
  T_emission em(open_HDU(file,"PARTICLEDATA"),true );

  cout << "setting up emergence" << endl;
  
  T_emergence e (3,5,10000,100,150);

  auto_ptr<FITS>  output;
  output.reset (new FITS ("galaxy.fits", Write));

  // load stuff
  //e.load_images(*output);
  //gr.load_intensity(open_HDU(*output,"INTENSITY"));
  seed(43);

  binifstream binf("dump");
  e.load_dump(binf);
  gr.load_dump(binf);

  cout << "shooting" << endl;

  const int n_threads = 1;
  const int n_rays=10000;
  std::vector<T_rng::T_state> states= generate_random_states (n_threads);
  T_rng_policy rng;

  xfer<T_dust_model, T_grid, T_rng_policy> x (gr, em, e, model, rng, 1e-4, 0.,1);
  long n_rays_actual = 0;
  dummy_terminator t;
  T_biaser b(300);
  for(;n_rays_actual<n_rays;n_rays_actual++) {
    T_biaser b(int(rnd()*n_lambda));
     //T_biaser b(300);
  x.shoot(b);
  }
  //shoot (x, scatter_shooter<T_biaser> (b), t, states, n_rays,
  //	 n_rays_actual, n_threads,0, "");

  T_unit_map units;
  units ["length"] = "kpc";
  units ["luminosity"] = "W";

  //binofstream binf("dump");
  //e.write_dump(binf);
  //gr.write_dump(binf);

  //e.write_parameters(output, units);
  //e->create_HDU(*output, 2, units);
  const T_float normalization = 1./n_rays_actual;
  //global_hdu = output.addImage ("junk", DOUBLE_IMG, n );
  //e->write_parameters(*output,units);

  e.write_images(*output,normalization,units,"",false,false);
  gr.write_intensity(*output, "INTENSITY", "DEPOSITION", 
		     normalization, model,false);

}

void aux_galaxy ()
{
  cout << "loading grid" << endl;

  FITS file("/data/sunrise/test/Sbc201a-u4/timetest/grid_012.fits",Read);

  // the aux grid is both a grid and an emission object
  typedef aux_grid<cumulative_sampling, local_random> T_grid;
  T_grid gr(open_HDU(file,"GRIDSTRUCTURE"),
	    open_HDU(file,"GRIDDATA"));

  cout << "setting up emergence" << endl;
  
  aux_grid_emergence e (3,5,10000,100,150);

  auto_ptr<FITS>  output;
  output.reset (new FITS ("galaxy.fits", Write));

  // load stuff
  //e.load_images(*output);
  //gr.load_intensity(open_HDU(*output,"INTENSITY"));
  seed(43);

  cout << "shooting" << endl;

  const int n_threads = 1;
  const int n_rays=10000000;
  std::vector<T_rng::T_state> states= generate_random_states (n_threads);
  T_rng_policy rng;

  typedef generic_dust_model<aux_pars_type> T_dust_model;
  T_dust_model model;

  xfer<T_dust_model, T_grid, T_rng_policy> x (gr, gr, e, model, rng, 1e-4, 0.,1);
  long n_rays_actual = 0;
  dummy_terminator t;
  T_dust_model::T_biaser b;
  for(;n_rays_actual<n_rays;n_rays_actual++) {
    x.shoot_isotropic(b);
  }
  //shoot (x, scatter_shooter<T_biaser> (b), t, states, n_rays,
  //	 n_rays_actual, n_threads,0, "");

  T_unit_map units;
  units ["length"] = "kpc";
  units ["luminosity"] = "W";

  //binofstream binf("dump");
  //e.write_dump(binf);
  //gr.write_dump(binf);

  //e.write_parameters(output, units);
  //e->create_HDU(*output, 2, units);
  const T_float normalization = 1./n_rays_actual;
  //global_hdu = output.addImage ("junk", DOUBLE_IMG, n );
  //e->write_parameters(*output,units);

  e.write_images(*output,normalization,units);
}

void aux_particle_galaxy ()
{
  cout << "loading emission" << endl;

  FITS file("/data/sunrise/test/Sbc201a-u4/timetest/grid_012.fits",Read);
  aux_particles_emission<cumulative_sampling, local_random> em(open_HDU(file,"PARTICLEDATA"));

  cout << "setting up emergence" << endl;

  aux_grid_emergence e (3,5,10000,100,150);

  auto_ptr<FITS>  output;
  output.reset (new FITS ("galaxy.fits", Write));

  // load stuff
  //e.load_images(*output);
  //gr.load_intensity(open_HDU(*output,"INTENSITY"));
  seed(43);

  cout << "shooting" << endl;

  const int n_threads = 1;
  const int n_rays=10000;
  std::vector<T_rng::T_state> states= generate_random_states (n_threads);
  T_rng_policy rng;

  ::dummy_grid gr;
  typedef generic_dust_model<aux_pars_type> T_dust_model;
  T_dust_model model;

  xfer<T_dust_model, ::dummy_grid, T_rng_policy> x (gr, em, e, model, rng, 1e-4, 0.,1);
  long n_rays_actual = 0;
  dummy_terminator t;
  T_dust_model::T_biaser b;
  for(;n_rays_actual<n_rays;n_rays_actual++) {
    x.shoot_isotropic(b);
  }
  //shoot (x, scatter_shooter<T_biaser> (b), t, states, n_rays,
  //	 n_rays_actual, n_threads,0, "");

  T_unit_map units;
  units ["length"] = "kpc";
  units ["luminosity"] = "W";

  //binofstream binf("dump");
  //e.write_dump(binf);
  //gr.write_dump(binf);

  //e.write_parameters(output, units);
  //e->create_HDU(*output, 2, units);
  const T_float normalization = 1./n_rays_actual;
  //global_hdu = output.addImage ("junk", DOUBLE_IMG, n );
  //e->write_parameters(*output,units);

  e.write_images(*output,normalization,units);
}

void intensity_test ()
{
  cout << "loading grid" << endl;
  const long n_rays =1000000;
  const int grid_n = 10;
  const int grid_ref = 0; 
  const T_float grid_extent=2;
  const int n_forced=1;
  T_densities rho(1); rho=1;
  seed(42);
  // k g a
  T_dust_model::T_scatterer_vector sv;
  const int n=7;
  T_lambda kk(n),gg(n),aa(n);
  kk=.1,.25,.5,1,2,4,8;
  gg=0;
  aa=0;
  sv.push_back(boost::shared_ptr<T_scatterer> (new simple_HG_dust_grain<polychromatic_scatterer_policy, T_rng_policy>(kk,gg,aa)));

  T_dust_model model (sv);
  T_lambda lambda(n);
  lambda = 0;//550e-9;
  model.set_wavelength(lambda);

  uniform_factory factory (rho,n, grid_ref);
  polychromatic_grid gr (vec3d (-1,-1,0),
			 vec3d (1,1,grid_extent),
			 vec3i (1, 1,grid_n), factory);

  cout << "setting up emission" << endl;

  typedef emission<polychromatic_policy, T_rng_policy> T_emission;
  auto_ptr<T_emission> em;
  T_lambda luminosity(n);
  luminosity=1;
  em.reset(new laser_emission<polychromatic_policy, T_rng_policy> (vec3d (.01, .01, -1), vec3d (0, 0, 1),luminosity));

  cout << "setting up emergence" << endl;
  
  std::vector<std::pair<T_float, T_float> > cam_pos;
  cam_pos.push_back(std:: make_pair (0.1, 0.01) );
  T_emergence e (cam_pos, 100, 1, 1);
  
  cout << "shooting" << endl;

  const int n_threads = 1;
  std::vector<T_rng::T_state> states= generate_random_states (n_threads);
  T_rng_policy rng;

  xfer<T_dust_model, polychromatic_grid, T_rng_policy> x (gr, *em, e, model, rng, 1e-4, 0,n_forced,10.);
  long n_rays_actual = 0;
  dummy_terminator t;
  //T_biaser b(3);
  //shoot (x, scatter_shooter<T_biaser> (b), t, states, n_rays,
  //	 n_rays_actual, n_threads,0, "");

  for(;n_rays_actual<n_rays;n_rays_actual++) {
    //T_biaser b(int(rnd()*n));
    T_biaser b(3);
    x.shoot(b);
  }

  const T_float normalization = 1.//em->total_emission_weight()
    /n_rays_actual;
  //global_hdu = output.addImage ("junk", DOUBLE_IMG, n );
  //e.write_images(output, emission.total_emission_weight()/N, 1);

  //FITS output ("junk.fits", Write);
  T_lambda total_absorbed_a(n);total_absorbed_a=0;
  T_lambda total_absorbed_i(n);total_absorbed_i=0;
  ofstream of("intens.txt");
  for(T_grid::iterator i = gr.begin(); i!=gr.end();++i) {
    total_absorbed_i+=normalization*i->data()->intensity()*kk*(1-aa);
    DEBUG(1,total_absorbed_a+=normalization*i->data()->absorbed(););
    of << i->getcenter()[2];
    for(int j=0;j<n;++j)
      of << '\t' << normalization*i->data()->intensity()(j)/(4*constants::pi*i->volume()) << '\t' << normalization*i->data()->intensity()(j)*kk(j)*(1-aa(j)) ;
	//<< '\t' << normalization*i->data()->absorbed()(j);
    of << endl;
  }
  cout << "Total absorbed luminosity by absorption: " << total_absorbed_a << endl;
  cout << "Total absorbed luminosity by intensity: " << total_absorbed_i << endl;
}

int main (int, char**)
{
  //w77();
  //aux_galaxy();
  intensity_test();
}


